import pandas as pd
import mysql.connector
from mysql.connector import Error


def fetch_data_from_mysql(host_name, user_name, user_password, db_name, query):
    """
    从MySQL数据库查询数据
    """
    connection = None
    try:
        connection = mysql.connector.connect(
            host=host_name,
            user=user_name,
            passwd=user_password,
            database=db_name
        )
        if connection.is_connected():
            db_Info = connection.get_server_info()
            print("Connected to MySQL Server version ", db_Info)
            cursor = connection.cursor()
            cursor.execute(query)
            records = cursor.fetchall()

            # 使用pandas的DataFrame来存储数据
            # 设置默认表头
            # columns = [col[0] for col in cursor.description]
            # 手动设置表头 `id`, `stat_date`, `stock_code`, `stock_name`, `stock_price`, `stock_price_pct`, `rzrqye`,`lrzl`,`rzjmr`, `rzche`,`rzmr`, `rzye`,`rqjmc`, `rqche`,`rqmc`, `rqyl`, `rqye`
            columns = ['id', '交易日期','股票代码','股票名称','股价','今日涨幅','融资融券余额','两融增量','融资净买入', '融资偿还额',  '融资买入', '融资余额', '融券净卖出', '融券偿还额','融券卖出', '融券余量','融券余额']
            df = pd.DataFrame(records, columns=columns)
            df['股价'] = pd.to_numeric(df['股价'], errors='coerce')
            df['今日涨幅'] = pd.to_numeric(df['今日涨幅'], errors='coerce')
            df['融资融券余额'] = (pd.to_numeric(df['融资融券余额'], errors='coerce')/10000).round(2)
            df['两融增量'] = (pd.to_numeric(df['两融增量'], errors='coerce')/10000).round(2)
            df['融资净买入'] = (pd.to_numeric(df['融资净买入'], errors='coerce') / 10000).round(2)
            df['融资偿还额'] = (pd.to_numeric(df['融资偿还额'], errors='coerce') / 10000).round(2)
            df['融资买入'] = (pd.to_numeric(df['融资买入'], errors='coerce')/10000).round(2)
            df['融资余额'] = (pd.to_numeric(df['融资余额'], errors='coerce')/10000).round(2)
            df['融券净卖出'] = (pd.to_numeric(df['融券净卖出'], errors='coerce')/10000).round(2)
            df['融券偿还额'] = (pd.to_numeric(df['融券偿还额'], errors='coerce')/10000).round(2)
            df['融券卖出'] = (pd.to_numeric(df['融券卖出'], errors='coerce')/10000).round(2)
            df['融券余量'] = (pd.to_numeric(df['融券余量'], errors='coerce')/10000).round(2)
            df['融券余额'] = (pd.to_numeric(df['融券余额'], errors='coerce')/10000).round(2)
            return df
    except Error as e:
        print("Error while connecting to MySQL", e)
    finally:
        if connection.is_connected():
            cursor.close()
            connection.close()
            print("MySQL connection is closed")


def export_to_excel(data, excel_path):
    """
    将数据导出到Excel文件
    """
    data.to_excel(excel_path, index=False)
    print(f"Data exported to {excel_path}")


# 示例用法
if __name__ == "__main__":
    host = "localhost"  # 数据库主机地址
    user = "root"  # 数据库用户名
    password = "123"  # 数据库密码
    database = "z_sproot_series"  # 数据库名称
    table = "stock_gnzj"  # 数据库表名

    query = """SELECT id, stat_date, stock_code, stock_name, stock_price, stock_price_pct, 
                rzrqye,lrzl,rzjmr, rzche,rzmr, rzye,rqjmc, rqche,rqmc, rqyl, rqye 
                FROM stock_rzrq a where stock_code = '300251' order by stat_date desc """
    excel_path = "output.xlsx"

    data = fetch_data_from_mysql(host, user, password, database, query)
    if data is not None:
        export_to_excel(data, excel_path)